/*
   \file  lazebnik.cc
   \brief Implementation of the lazebnik algorithm.

   This file defines the non-inline member functions of the
   gist::lazebnik class. See the header file for more info.
*/

/*
   This file is part of libgist.

   libgist is free software; you can redistribute it and/or modify it
   under the terms of the GNU General Public License as published by the
   Free Software Foundation; either version 2 of the License, or (at your
   option) any later version.

   libgist is distributed in the hope that it will be useful, but WITHOUT
   ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
   FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
   for more details.

   You should have received a copy of the GNU General Public License
   along with libgist; if not, write to the Free Software Foundation,
   Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA.
*/

/*
   REVISION HISTORY

   $HeadURL: http://libgist.googlecode.com/svn/branches/siagian_dev/src/lazebnik.cc $
   $Id: lazebnik.cc 113 2012-01-22 05:06:06Z libgist@gmail.com $
*/

//------------------------------ HEADERS --------------------------------

// libgist
#include "lazebnik.hh"
#include "kmeans.hh"
#include "threaded_image_set_processor.hh"
#include "settings.hh"
#include "image_utils.hh"
#include "local_fn.hh"

// Boost.Thread
#include <boost/thread/locks.hpp>
#include <boost/thread/mutex.hpp>

// Boost.Lambda
#include <boost/lambda/if.hpp>
#include <boost/lambda/bind.hpp>
#include <boost/lambda/lambda.hpp>

// Boost.Timer
#include <boost/timer.hpp>

// Boost.Foreach
#include <boost/foreach.hpp>
#define foreach BOOST_FOREACH

// Boost.uBLAS
#include <boost/numeric/ublas/matrix_proxy.hpp>
#include <boost/numeric/ublas/matrix.hpp>

// Boost.MultiArray
#include <boost/multi_array.hpp>

// Boost.Math
#include <boost/math/constants/constants.hpp>

// Standard C++
#include <iostream>
#include <numeric>
#include <algorithm>
#include <map>
#include <list>
#include <vector>
#include <string>
#include <iterator>
#include <utility>

// Standard C
#include <math.h>

// Unix
#include <sys/time.h>

//----------------------- NAMESPACE DEFINITION --------------------------

namespace gist {

//----------------------------- CONSTANTS -------------------------------

namespace {

// Gaussians play a large role in the SIFT algorithm. This setting
// specifies the standard deviation to use for these Gaussians.
const float SIGMA = 1.6f ;

// SIFT descriptors are computed in pixel patches centered at a number of
// keypoints. This parameter specifies the half-size of these patches.
//
// NOTE: This constant is called RADIUS because SIFT descriptor
// computation also involves weighted gradient magnitudes that drop off
// normally (i.e., Gaussian) from the keypoint center in proportion to
// the radial distance between a pixel in the patch and the center.
const int RADIUS = static_cast<int>(5 * SIGMA + .5f) ;

} // end of local anonymous namespace encapsulating above constants

//---------------------- SIFT DESCRIPTOR HELPERS ------------------------

// Quick helper to wrap around the notion of a SIFT descriptor, which is
// simply a vector of 128 numbers.
namespace {

struct sift_descriptor {
   static const int SIZE = 128 ;

   // Initialize an empty SIFT descriptor, i.e., all zeros.
   sift_descriptor() ;

   // Return reference to i-th element of the descriptor.
   float& operator[](int i) {return m_values[i] ;}
   const float& operator[](int i) const {return m_values[i] ;}

   // Iterator interface to descriptor values.
   float* begin() {return m_values ;}
   float* end  () {return m_values + SIZE ;}
   const float* begin() const {return m_values ;}
   const float* end  () const {return m_values + SIZE ;}

   // Check if the entire descriptor consists of zeros.
   bool is_zero() const ;

private :
   float m_values[SIZE] ;
} ;

// SIFT descriptor initialization
sift_descriptor::sift_descriptor()
{
   const int n = sizeof m_values / sizeof m_values[0] ;
   std::fill_n(m_values, n, 0) ;
}

// Check if SIFT descriptor is all zeros
bool sift_descriptor::is_zero() const
{
   for (int i = 0; i < SIZE; ++i)
      if (m_values[i] < -1e-6f || m_values[i] > 1e-6f) // non-zero
         return false ;
   return true ;
}

// This class acts as a temporary store for SIFT keypoint features. The
// final feature vector in a SIFT keypoint is an array of bytes. However,
// here, we use floats to allow normalizations and other such operations.
//
// As detailed in David Lowe's 2004 SIFT paper, a SIFT descriptor is
// essentially a histogram with 128 bins (4 bins for x * 4 bins for y * 8
// bins for orientation).
class feature_vector {
   // The sizes of the different bins.
   static const int X_SIZE = 4 ;
   static const int Y_SIZE = 4 ;
   static const int Z_SIZE = 8 ;

   // Some precomputed values to try and speed things up.
   static const float X_SIZE_BOUNDARY = X_SIZE - 0.5f ;
   static const float Y_SIZE_BOUNDARY = Y_SIZE - 0.5f ;
   static const float Z_SIZE_BOUNDARY = Z_SIZE - 0.5f ;

   static const int ZY_SIZE = Z_SIZE * Y_SIZE ;

   // As detailed in David Lowe's 2004 SIFT paper, the descriptor values
   // are normalized, thresholded and then renormalized. This parameter
   // specifies the threshold limit to use.
   //
   // DEVNOTE: For some odd reason, debug builds fail when this constant
   // is defined inside the class. On 32-bit platforms, the linker
   // complains about R_386_GOTOFF; and on 64-bit platforms, it complains
   // about R_X86_64_PC32. In both cases, the problem has something to do
   // with relocation against an undefined symbol (viz., THRESHOLD) not
   // being allowed when making a shared object.
   //
   // Apparently, it's okay when building a static library instead of a
   // .so. It is also okay in release builds!
   //
   // The workaround (that works for release, debug, static and shared
   // libraries) is to define this constant outside the class. Not sure
   // what the hell is going on here...
   //
   //static const float THRESHOLD = 0.2f ;
   static const float THRESHOLD ; // define outside

   // The SIFT features are stored in an array of floats whose size
   // equals the total number of bins.
   std::vector<float> m_features ;

public:
   // Create empty bins (i.e., initialized to zero)
   feature_vector() ;

   // Add the specified value to appropriate bins of SIFT descriptor
   void add(float x, float y, float z, float value) ;

   // Convert the features to a SIFT descriptor
   sift_descriptor to_sift() const ;

private:
   // Normalize array using sum of squares rather than just the sum
   std::vector<float> normalize(std::vector<float>) const ;

   // Threshold given array so that all values are capped to THRESHOLD
   std::vector<float> threshold(std::vector<float>) const ;

   // Clamp feature value to lie in [0, 255]
   float clamp(float) const ;
} ;

// Debug shared library can't be built when this is inside class. Go figure...
const float feature_vector::THRESHOLD = .2f ;

feature_vector::feature_vector()
   : m_features(X_SIZE * Y_SIZE * Z_SIZE, 0)
{}

void feature_vector::add(float x, float y, float z, float value)
{
   int   xi0, xi1, yi0, yi1, zi0, zi1 ; // the bins
   float wx0, wy0, wz0, wx1, wy1, wz1 ; // the corresponding weights

   // If close to bounds then the values go fully into the end bins,
   // otherwise they split between two adjacent bins.
   // Note: a value of 2.0 should equally split between bins 1 and 2.
   if (x <= 0.5f)
   {
      xi0 = 0 ;
      xi1 = 0 ;
      wx0 = 0.5f ;
      wx1 = 0.5f ;
   }
   else if (x >= X_SIZE_BOUNDARY)
   {
      xi0 = X_SIZE - 1 ;
      xi1 = X_SIZE - 1 ;
      wx0 = 0.5f ;
      wx1 = 0.5f ;
   }
   else
   {
      const float xx = x - 0.5f ;
      xi0 = static_cast<int>(xx);
      xi1 =  1 + xi0 ;
      wx1 = xx - xi0 ;
      wx0 =  1 - wx1 ;
   }
   if (y <= 0.5f)
   {
      yi0 = 0 ;
      yi1 = 0 ;
      wy0 = 0.5f ;
      wy1 = 0.5f ;
   }
   else if (y >= Y_SIZE_BOUNDARY)
   {
      yi0 = Y_SIZE - 1 ;
      yi1 = Y_SIZE - 1 ;
      wy0 = 0.5f ;
      wy1 = 0.5f ;
   }
   else
   {
      const float yy = y - 0.5f ;
      yi0 = static_cast<int>(yy);
      yi1 =  1 + yi0 ;
      wy1 = yy - yi0 ;
      wy0 =  1 - wy1 ;
   }
   // The situation is slightly different for orientation (i.e., z)
   // because we wrap around.
   if (z <= 0.5f)
   {
      zi0 = 0 ;
      zi1 = Z_SIZE - 1 ;
      wz0 =   0.5f + z ;
      wz1 =   1.0f - wz0 ;
   }
   else if (z >= Z_SIZE_BOUNDARY)
   {
      zi0 = Z_SIZE - 1 ;
      zi1 = 0 ;
      wz0 = Z_SIZE + 0.5f - z ;
      wz1 = 1 - wz0 ;
   }
   else
   {
      const float zz = z - 0.5f ;
      zi0 = static_cast<int>(zz);
      zi1 =  1 + zi0 ;
      wz1 = zz - zi0 ;
      wz0 =  1 - wz1 ;
   }

   // Convention: we add 1 for each unit of orientation (our fastest
   // varying index), then zSize for each unit of y, finally zSize*ySize
   // for each unit of x.
   xi0 *= ZY_SIZE ;
   xi1 *= ZY_SIZE ;
   yi0 *= Z_SIZE ;
   yi1 *= Z_SIZE ;

   m_features[xi0 + yi0 + zi0] += value * wx0 * wy0 * wz0 ;
   m_features[xi1 + yi0 + zi0] += value * wx1 * wy0 * wz0 ;
   m_features[xi0 + yi1 + zi0] += value * wx0 * wy1 * wz0 ;
   m_features[xi1 + yi1 + zi0] += value * wx1 * wy1 * wz0 ;
   m_features[xi0 + yi0 + zi1] += value * wx0 * wy0 * wz1 ;
   m_features[xi1 + yi0 + zi1] += value * wx1 * wy0 * wz1 ;
   m_features[xi0 + yi1 + zi1] += value * wx0 * wy1 * wz1 ;
   m_features[xi1 + yi1 + zi1] += value * wx1 * wy1 * wz1 ;
}

sift_descriptor feature_vector::to_sift() const
{
   using boost::lambda::_1 ;
   using boost::lambda::bind ;

   sift_descriptor s ;
   std::vector<float> f = normalize(threshold(normalize(m_features))) ;
   std::transform(f.begin(), f.end(), s.begin(),
                  bind(&feature_vector::clamp, this, _1 * 512)) ;
   return s ;
}

std::vector<float> feature_vector::normalize(std::vector<float> v) const
{
   using boost::lambda::_1 ;
   using boost::lambda::_2 ;

   // DEVNOTE: std::accumulate() invokes the function object passed as
   // its last argument like so: f(a, b), where 'a' is the third argument
   // and 'b' the elements of the sequence. Thus, in terms of
   // Boost.Lambda's placeholders, 'a' corresponds to _1 and 'b' to _2,
   // which is different from most use-cases involving the STL algorithms
   // and Boost.Lambda placeholders where _1 usually corresponds to the
   // container elements.
   float sum = std::accumulate(v.begin(), v.end(), 0.0f, _1 + _2*_2) ;
   if (sum < 1e-10f) // sum of squares too small to normalize
      return v ;

   float normalizer = 1/sqrtf(sum) ;
   std::transform(v.begin(), v.end(), v.begin(), _1*normalizer) ;
   return v ;
}

std::vector<float> feature_vector::threshold(std::vector<float> v) const
{
   using boost::lambda::_1 ;
   using boost::lambda::if_then_else ;
   std::for_each(v.begin(), v.end(),
                 if_then_else(_1 > THRESHOLD, _1 = THRESHOLD, _1 = _1)) ;
   return v ;
}

float feature_vector::clamp(float f) const
{
   if (f < 0)
      return 0 ;
   if (f > 255)
      return 255 ;
   return f ;
}

// This function computes the SIFT descriptor of a 16x16 pixel patch
// centered at the point (x,y).
sift_descriptor
compute_sift_descriptor(int x, int y, const gradient_image_t& grad)
{
   // Just some precomputed constants that allow us to minimize floating
   // point operations inside this function's inner loop.
   const float pi = boost::math::constants::pi<float>() ;
   const float _1_over_2sigma_sqr = -1/(2.0f * RADIUS * RADIUS) ;
   const float _2pi               = 2*pi ;
   const float _1_over_pi_by_4    = 4/pi ;
   const float _2_over_r          = 2.0f/RADIUS ;

   const int w = grad.width() ;
   const int h = grad.height();

   // Setup a "pointer" to the gradient pixel located at (x,y). To get at
   // the surrounding pixels in the patch centered at (x,y), we will
   // apply an offset from this location.
   typedef gradient_image_t::const_view_t view ;
   view v = boost::gil::const_view(grad) ;
   view::xy_locator pixel = v.xy_at(x,y) ;

   // Iterate over a square pixel patch of size 2*RADIUS + 1 centered at
   // (x,y) and extract the SIFT features in this area of the image by
   // binning the gradients therein and weighting the bins using a
   // Gaussian that drops off according to the radial distance between
   // each point in the patch and patch's center (x,y).
   feature_vector fv ;
   for (int j = -RADIUS; j < RADIUS; ++j)
   for (int i = -RADIUS; i < RADIUS; ++i)
   {
      const int img_x = x + i ;
      const int img_y = y + j ;
      if (img_x < 0 || img_x >= w || img_y < 0 || img_y >= h)
         continue ; // image coordinates out of bounds; try next pixel...

      // Find the dominant orientation at (i,j). This is the SIFT feature
      // for this pixel. This orientation is assigned to one (or more) of
      // eight possible bins (i.e., 0-45, 45-90, etc.).
      //
      // NOTE: All the range checking with pi is just to bring the angle
      // into the [0, 360] range.
      const gradient_pixel_t& gradient_vector = pixel(i,j) ;
      float gradient_angle = fmodf(boost::gil::at_c<1>(gradient_vector), _2pi) ;
      if (gradient_angle < 0)
         gradient_angle += _2pi ;
      if (gradient_angle >= pi)
         gradient_angle -= _2pi ;
      const float orientation = // dominant orientation split into 8 bins
         4 + _1_over_pi_by_4 * gradient_angle ;

      // The feature at (i,j) will be normally (i.e., Gaussian) weighted
      // based on the distance between (i,j) and (x,y).
      const float exp = expf((i*i + j*j) * _1_over_2sigma_sqr) ;
      const float weight = exp * boost::gil::at_c<0>(gradient_vector) ;

      // Add the feature to the feature vector at fractional coordinates
      // corresponding to a 4x4 grid of bins.
      const float xf = 2 + _2_over_r * i ;
      const float yf = 2 + _2_over_r * j ;
      fv.add(xf, yf, orientation, weight);
   }
   return fv.to_sift() ;
}

} // end of local anonymous namespace encapsulating above helper

//----------------------- SIFT DESCRIPTOR GRIDS -------------------------

namespace {

// The Lazebnik algorithm works by computing SIFT descriptors for the
// input image arranged in a regular grid (no SIFT keypoints, etc.).
// This class is a thin wrapper around a 2D array for holding an image's
// SIFT descriptors arranged in a grid.
class sift_grid {
   // This is the data structure that holds the SIFT descriptors for an
   // image.
   boost::multi_array<sift_descriptor, 2> m_grid ;

   // Some parts of images can have SIFT descriptors that are entirely
   // zero (e.g., a 16x16 pixel patch that has no or very little
   // texture). We do not want to use these zero descriptors for
   // training the classifier. Therefore, as we add descriptors to the
   // grid, it is useful to keep track of the number of "valid" (i.e.,
   // non-zero) descriptors the grid contains.
   int m_num_nonzero_descriptors ;

public:
   // Create a grid of the specified width and height
   sift_grid(int w, int h) ;

   // Add a SIFT descriptor to the specified cell of the grid and update
   // the non-zero descriptors count as appropriate.
   void add(int i, int j, const sift_descriptor& d) ;

   // Return the number of non-zero descriptors in the grid. This method
   // is meant to be called after all the SIFT descriptors for a
   // particular grid have been added.
   int num_nonzero_descriptors() const {return m_num_nonzero_descriptors ;}

   // This type facilitates easy iteration over each of the SIFT
   // descriptors in the grid.
   typedef boost::multi_array_ref<const sift_descriptor, 1> grid_data ;

   // This method returns an instance of the above type to allow clients
   // to iterate over each of the SIFT descriptors contained by this
   // grid.
   grid_data data() const {
      return grid_data(m_grid.data(), boost::extents[m_grid.num_elements()]) ;
   }
} ;

// Constructor: create 2D array of specified size
sift_grid::sift_grid(int w, int h)
   : m_grid(boost::extents[w][h]), m_num_nonzero_descriptors(0)
{}

// Add a SIFT descriptor to the grid and update non-zero descriptor count
void sift_grid::add(int i, int j, const sift_descriptor& d)
{
   m_grid[i][j] = d ;
   if (!d.is_zero())
      ++m_num_nonzero_descriptors ;
}

// Given an input image, compute its grid of SIFT descriptors
sift_grid compute_sift_grid(const boost::gil::gray8_image_t& img)
{
   gradient_image_t grad = gradient(img) ;

   const int w = img.width() ;
   const int h = img.height();

   sift_grid g(w/RADIUS - 1, h/RADIUS - 1) ;
   for (int y = RADIUS, j = 0; y < h - RADIUS + 1; y += RADIUS, ++j)
   for (int x = RADIUS, i = 0; x < w - RADIUS + 1; x += RADIUS, ++i)
      g.add(i, j, compute_sift_descriptor(x, y, grad)) ;

   return g ;
}

} // end of local anonymous namespace encapsulating above helpers

//-------------------------- SIFT GRID LISTS ----------------------------

namespace {

// This data structure holds all the SIFT grids for a particular image
// category. It is used in training mode.
typedef std::list<sift_grid> sift_grid_list ;

// Quick helper to create a new list of SIFT descriptor grids containing
// the grid passed in to the function.
inline sift_grid_list new_list(const sift_grid& g)
{
   sift_grid_list list ;
   list.push_back(g) ;
   return list ;
}

} // end of local anonymous namespace encapsulating above helpers

//---------------------- LAZEBNIK IMPLEMENTATION ------------------------

// The lazebnik class hides its entire implementation behind an opaque
// object. This section defines this implementation class.
namespace {

// In training mode, we will accumulate the SIFT descriptors for the
// training images in a map of lists of SIFT grids indexed by the
// category labels.
typedef std::map<std::string, sift_grid_list> sift_grid_map ;

// Part of the training involves k-means clustering, which needs a matrix
// as its input.
typedef boost::numeric::ublas::matrix<float> matrix ;

// This class actually implements the Lazebnik algorithm. The interface
// class defined in the header file, for the most part, simply forwards
// its API to this one.
class lazebnik_impl {
   // This data structure holds grids of SIFT descriptors for the input
   // images arranged according to their categories.
   sift_grid_map m_map ;

   // A variable to keep track of the total number of SIFT descriptors in
   // the above map.
   int m_num_descriptors ;

   // As we allow multithreaded clients, we will have to synchronize
   // access to the above members.
   mutable boost::mutex m_mutex ;

public:
   // Constructor
   lazebnik_impl() ;

   // Add a new SIFT grid given its category
   void add(const std::string& cat, const sift_grid&) ;

   // Copy all categorized SIFT grids to a flat matrix
   matrix sift_grid_map_to_matrix() const ;
} ;

// Constructor
lazebnik_impl::lazebnik_impl()
   : m_num_descriptors(0)
{}

// Lock map of SIFT descriptor grids and add new grid given its category
void lazebnik_impl::add(const std::string& cat, const sift_grid& grid)
{
   boost::lock_guard<boost::mutex> lock(m_mutex) ;
   sift_grid_map::iterator m = m_map.find(cat) ;
   if (m == m_map.end())
      m_map.insert(std::make_pair(cat, new_list(grid))) ;
   else
      m->second.push_back(grid) ;
   m_num_descriptors += grid.num_nonzero_descriptors() ;
}

// Convert map of SIFT descriptor grids to a d-by-n matrix where d is
// the number of dimensions in a SIFT descriptor (i.e., 128) and n is
// the total number of non-zero SIFT descriptors stored in all the grids
// across all the categories.
matrix lazebnik_impl::sift_grid_map_to_matrix() const
{
   // Local helper class to copy a SIFT grid's descriptors as column
   // vectors in a matrix.
   class copy: public binary_fn<void, std::string, sift_grid> {
      matrix& m_matrix ;            // destination matrix
      mutable boost::mutex m_mutex ;// multiple threads are used to copy
      mutable int m_column ;        // index where descriptors should be copied
   public:
      copy(matrix& m): m_matrix(m), m_column(0) {}

      // Thread-safe copying of descriptors from given SIFT grid to
      // destination matrix. Each thread copies all the descriptors from
      // one grid into a block of columns in the matrix. Thus, only
      // the column index computation part is a critical section. The
      // copying itself can be outside the critical section, which should
      // considerably speed up the copy operation on a multicore system.
      //
      // NOTE: First parameter is image category label, which we do not
      // care about here.
      void operator()(const std::string&, const sift_grid& g) const {
         int j ; // start column index where descriptors in g should be copied

         // Local block to automatically acquire and release mutex
         {
            boost::lock_guard<boost::mutex> lock(m_mutex) ;
            j = m_column ;
            m_column += g.num_nonzero_descriptors() ;
         }


         // Copy descriptors to matrix columns; no need for locking...
         typedef boost::numeric::ublas::matrix_column<matrix> column_vector ;
         foreach(const sift_descriptor& descr, g.data())
            if (!descr.is_zero())
               std::copy(descr.begin(), descr.end(),
                         column_vector(m_matrix, j++).begin()) ;
      }
   } ;

   // Create a matrix large enough to accommodate all the non-zero
   // descriptors collected in all the grids across all categories.
   // Since this can be a fairly large matrix, use multiple threads to
   // copy the descriptors into the matrix.
   matrix m(sift_descriptor::SIZE, m_num_descriptors) ;
   copy copy_sift_grid_to_matrix(m) ;
   threaded_image_set_processor<sift_grid_list>(m_map).
      run(num_threads(), local_fn(copy_sift_grid_to_matrix)) ;
   return m ;
}

} // end of local anonymous namespace encapsulating above implementation details

//------------------------ LAZEBNIK INTERFACE ---------------------------

// Initialization
lazebnik::lazebnik()
   : p_impl(new lazebnik_impl)
{}

// Training
void
lazebnik::
data(const std::string& cat, const boost::gil::gray8_image_t& img)
{
   reinterpret_cast<lazebnik_impl*>(p_impl)->add(cat, compute_sift_grid(img)) ;
}

void lazebnik::train()
{
   lazebnik_impl* pimpl = reinterpret_cast<lazebnik_impl*>(p_impl) ;

   matrix m = pimpl->sift_grid_map_to_matrix() ;
   boost::timer t ;
   matrix vocabulary = kmeans<>(num_threads()).centroids(200, m) ;
   double s = t.elapsed()/num_threads() ;
   std::cout << "k-means clustering took " << s << " seconds\n" ;
}

void lazebnik::save_vocabulary(const std::string&) const
{
}

void lazebnik::save_classifier(const std::string&) const
{
}

// Clean-up
lazebnik::~lazebnik()
{
   delete reinterpret_cast<lazebnik_impl*>(p_impl) ;
}

//-----------------------------------------------------------------------

} // end of namespace encapsulating this file's definitions

/* So things look consistent in everyone's emacs... */
/* Local Variables: */
/* indent-tabs-mode: nil */
/* End: */
